package org.breathe.dynamic.thread.pool.core.ai.memory;

import dev.langchain4j.data.message.AiMessage;
import dev.langchain4j.data.message.ChatMessage;
import dev.langchain4j.data.message.UserMessage;
import dev.langchain4j.memory.ChatMemory;
import org.springframework.data.redis.core.StringRedisTemplate;

import java.util.List;

/**
 * @author: breathe
 * @createTime: 2025-06-27
 */
public class RedisChatMemory implements ChatMemory {

    private final String memoryId;
    private final String key;
    private final StringRedisTemplate redisTemplate;
    private final int windowSize;

    public RedisChatMemory(String memoryId, StringRedisTemplate redisTemplate, int windowSize) {
        this.memoryId = memoryId;
        this.redisTemplate = redisTemplate;
        this.windowSize = windowSize;
        this.key = "chat:memory:" + memoryId;
    }

    @Override
    public Object id() {
        return memoryId;
    }

    @Override
    public void add(ChatMessage message) {
        redisTemplate.opsForList().rightPush(key, serialize(message));
        redisTemplate.opsForList().trim(key, -windowSize, -1); // 保留最新 N 条
    }

    @Override
    public List<ChatMessage> messages() {
        List<String> raw = redisTemplate.opsForList().range(key, 0, -1);
        if (raw == null) return List.of();
        return raw.stream().map(this::deserialize).toList();
    }

    @Override
    public void clear() {
        redisTemplate.delete(key);
    }

    private String serialize(ChatMessage message) {
        String role;
        String text;
        if (message instanceof UserMessage) {
            role = "USER";
            text = ((UserMessage) message).singleText();
        } else if (message instanceof AiMessage) {
            role = "AI";
            text = ((AiMessage) message).text();
        } else {
            throw new IllegalArgumentException("Unsupported message type: " + message.getClass());
        }
        return role + "|" + text.replace("\n", "\\n");
    }

    private ChatMessage deserialize(String raw) {
        int splitIndex = raw.indexOf('|');
        String role = raw.substring(0, splitIndex);
        String text = raw.substring(splitIndex + 1).replace("\\n", "\n");

        return switch (role) {
            case "USER" -> new UserMessage(text);
            case "AI" -> new AiMessage(text);
            default -> throw new IllegalArgumentException("Unknown role: " + role);
        };
    }

}
